#include "kpextOpImgCpyF32C1.h"
#include "kpextSpvData.h"

using namespace kpext;

OpImgCpyF32C1::OpImgCpyF32C1(kpextVkContext *ctx) :
    BaseCompute(ctx)
{
    if(!m_vkCtx){
        throw std::runtime_error("OpImgCpyF32C1 construct with null pointer");
    }

    m_inputCols = 1;
    m_inputRows = 1;
    m_outputCols = 1;
    m_outputRows = 1;

    m_postBarrier = true;

    //  初始化描述符集及布局
    m_descrPool = m_vkCtx->descriptorPool(
                {2},
                1);
    vk::DescriptorSetLayoutBinding binding(
                0,
                vk::DescriptorType::eStorageBuffer,
                1,
                vk::ShaderStageFlagBits::eCompute
                );
    std::vector<vk::DescriptorSetLayoutBinding> bindings;
    bindings.push_back(
                binding
                );
    binding.setBinding(1);
    bindings.push_back(binding);

    m_descrSetlayout = m_descrPool->createDescriptorSetLayout(
                bindings
                );

    m_descrSet = m_descrPool->recreateDescSet(*m_descrSetlayout,
                                              m_descrSet
                                              );


    // 初始化pipeline
    m_pipeline = m_vkCtx->computePipeline();
    m_pipeline->m_specConsts = std::make_shared<PipelineConstSZ>(3, 4);
    m_pipeline->m_pushConsts = std::make_shared<PipelineConstSZ>(4, 4);

    m_workGroupSize = ctx->getDefWorkGroupSize();

    m_pipeline->m_specConsts->set({m_workGroupSize[0],
                                   m_workGroupSize[1],
                                   1
                                  });
    m_pipeline->m_pushConsts->set({1, 1, 1, 1});

    m_pipeline->m_descrSetLayouts.push_back(m_descrSetlayout);

    if(!m_pipeline->rebuild(
                sizeof (spirv::ImgCpyF32C1), spirv::ImgCpyF32C1)
            ){
        SPDLOG_WARN("Rebuild kpext::ComputePipeline failed !");
    }

    //  工作组数量
    m_workCount = calc2DWorkGroup(1, 1,
                                  m_workGroupSize);
}

OpImgCpyF32C1::~OpImgCpyF32C1()
{

}

void OpImgCpyF32C1::rebuild(int inputCols,
                            int inputRows,
                            int outputCols,
                            int outputRows,
                            std::shared_ptr<BaseBuffer> inbuf,
                            std::shared_ptr<BaseBuffer> outbuf)
{
    if(!inbuf || !outbuf || !m_vkCtx){
        throw std::runtime_error(__FUNCTION__ " input with null pointer!");
    }

    if(inputCols < 0) inputCols = 1;
    if(inputRows < 0) inputRows = 1;
    if(outputCols < 0) outputCols = 1;
    if(outputRows < 0) outputRows = 1;


    m_cmdBarrier.prepare(
                {outbuf,},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eShaderRead
                );

    STCAST_ASSIGN(m_inputCols, inputCols);
    STCAST_ASSIGN(m_inputRows, inputRows);
    STCAST_ASSIGN(m_outputCols, outputCols);
    STCAST_ASSIGN(m_outputRows, outputRows);

    if(m_inBuf != inbuf.get()
            || m_outBuf != outbuf.get()
            ){
        BaseCompute::rebuild({inbuf, outbuf});
        m_inBuf = inbuf.get();
        m_outBuf = outbuf.get();
    }

    //  设定输入输出尺寸
    m_pipeline->m_pushConsts->set(
                {m_inputCols, m_inputRows,
                 m_outputCols, m_outputRows
                }
                );


    //  重新计算工作组大小
    m_workCount = calc2DWorkGroup(m_outputCols, m_outputRows,
                                  m_workGroupSize);

}
